from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus
from pytorch_grad_cam.utils.image import show_cam_on_image
import torch
import torch.nn as nn
from Utils import get_awa2_images
import cv2
from LoadModel import load_alexnet, load_vggnet, load_resnet
import sys
sys.path.append("./")


# model = load_vggnet(50, 3)
model = load_resnet(50, 3)
model.load_state_dict(torch.load("./Model/AwA2/resnet_model_epoch399.pth"))
print(model)
# exit(0)
# layer = model.conv[17]
# layer = model.conv[4][5]
layer = model.layer4[1].left[3]
print(layer)
# exit(0)

cam = GradCAM(model=model, target_layer=layer, use_cuda=True)
origin_images, input_images, labels = get_awa2_images(128, 224)
grayscale_cam = cam(input_tensor=input_images, target_category=labels, aug_smooth=True)
grayscale_cam = grayscale_cam[100, :]
# exit(0)
visualization = show_cam_on_image(origin_images[20, :], grayscale_cam, use_rgb=False)
print(visualization.shape)
print(type(visualization))
cv2.imwrite("testcam.png", visualization)